import numpy as np
from sklearn.cluster import KMeans
from PIL import ImageGrab
import time
import logging

logger = logging.getLogger(__name__)

class KMeansColorDetector:
    """基于K-Means聚类的颜色检测器"""

    def __init__(self, n_clusters=3, tolerance=30):
        self.n_clusters = n_clusters
        self.tolerance = tolerance
        self.kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)

        # 性能统计
        self.detection_times = []
        self.cluster_cache = {}
        self.cache_max_age = 1.0  # 缓存1秒

    def detect_dominant_colors(self, image_region):
        """检测图像区域的主要颜色"""
        try:
            start_time = time.time()

            # 将图像转换为像素数组
            if hasattr(image_region, 'getdata'):
                pixels = list(image_region.getdata())
            else:
                pixels = image_region

            if not pixels:
                return None, []

            # 转换为numpy数组
            pixel_array = np.array(pixels).reshape(-1, 3)

            # 检查唯一颜色数量
            unique_colors = np.unique(pixel_array, axis=0)
            actual_clusters = min(self.n_clusters, len(unique_colors))

            # 如果只有一种颜色，直接返回
            if len(unique_colors) == 1:
                dominant_colors = unique_colors
                weights = np.array([1.0])
                detection_time = time.time() - start_time
                self.detection_times.append(detection_time)
                logger.debug(f"单一颜色检测: {dominant_colors[0]}, 耗时: {detection_time:.3f}s")
                return dominant_colors, weights

            # 动态调整聚类数量
            if actual_clusters != self.n_clusters:
                logger.debug(f"调整聚类数量: {self.n_clusters} -> {actual_clusters}")
                kmeans = KMeans(n_clusters=actual_clusters, random_state=42, n_init=10)
            else:
                kmeans = self.kmeans

            # 执行K-Means聚类，抑制警告
            try:
                from .warning_suppressor import suppress_kmeans_warnings
                with suppress_kmeans_warnings():
                    kmeans.fit(pixel_array)
            except ImportError:
                # 回退到基本警告抑制
                import warnings
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore", category=UserWarning)
                    warnings.filterwarnings("ignore", message=".*Number of distinct clusters.*")
                    try:
                        from sklearn.exceptions import ConvergenceWarning
                        warnings.filterwarnings("ignore", category=ConvergenceWarning)
                    except ImportError:
                        pass
                    kmeans.fit(pixel_array)

            # 获取聚类中心（主要颜色）
            cluster_centers = kmeans.cluster_centers_.astype(int)

            # 计算每个聚类的权重（像素数量）
            labels = kmeans.labels_
            unique_labels, counts = np.unique(labels, return_counts=True)

            # 按权重排序
            sorted_indices = np.argsort(counts)[::-1]
            dominant_colors = cluster_centers[sorted_indices]
            weights = counts[sorted_indices] / len(pixels)

            detection_time = time.time() - start_time
            self.detection_times.append(detection_time)

            logger.debug(f"K-Means检测耗时: {detection_time:.3f}s, 实际聚类数: {len(dominant_colors)}")

            return dominant_colors, weights

        except Exception as e:
            logger.error(f"K-Means颜色检测失败: {e}")
            return None, []

    def is_color_match_kmeans(self, target_color, image_region, match_threshold=0.3):
        """使用K-Means聚类进行颜色匹配"""
        try:
            dominant_colors, weights = self.detect_dominant_colors(image_region)

            if dominant_colors is None:
                return False, 0.0

            target_rgb = np.array(target_color)
            best_match_score = 0.0

            # 检查每个主要颜色是否匹配目标颜色
            for i, color in enumerate(dominant_colors):
                # 计算颜色距离
                distance = np.linalg.norm(target_rgb - color)

                if distance <= self.tolerance:
                    # 匹配成功，计算匹配分数（考虑权重）
                    match_score = weights[i] * (1 - distance / self.tolerance)
                    best_match_score = max(best_match_score, match_score)

            is_match = best_match_score >= match_threshold

            logger.debug(f"K-Means匹配结果: {is_match}, 分数: {best_match_score:.3f}")

            return is_match, best_match_score

        except Exception as e:
            logger.error(f"K-Means颜色匹配失败: {e}")
            return False, 0.0

    def get_performance_stats(self):
        """获取性能统计"""
        if not self.detection_times:
            return {"avg_time": 0, "total_detections": 0}

        return {
            "avg_time": np.mean(self.detection_times),
            "min_time": np.min(self.detection_times),
            "max_time": np.max(self.detection_times),
            "total_detections": len(self.detection_times)
        }

class HybridColorDetector:
    """混合颜色检测器 - 结合传统方法和K-Means"""

    def __init__(self):
        self.kmeans_detector = KMeansColorDetector(n_clusters=3)
        self.use_kmeans_threshold = 100  # 像素数量超过此值时使用K-Means

    def detect_color_hybrid(self, target_color, image_region, algorithm="RGB",
                          tolerance=10, hsv_tolerances=None):
        """混合颜色检测方法"""
        try:
            # 获取像素数据
            if hasattr(image_region, 'getdata'):
                pixels = list(image_region.getdata())
            else:
                pixels = image_region

            pixel_count = len(pixels)

            # 根据像素数量选择检测方法
            if pixel_count > self.use_kmeans_threshold:
                # 大区域使用K-Means聚类
                logger.debug(f"使用K-Means检测 (像素数: {pixel_count})")
                is_match, confidence = self.kmeans_detector.is_color_match_kmeans(
                    target_color, pixels
                )
                return is_match, confidence, "kmeans"
            else:
                # 小区域使用传统方法
                logger.debug(f"使用传统方法检测 (像素数: {pixel_count})")
                is_match, confidence = self._traditional_detection(
                    target_color, pixels, algorithm, tolerance, hsv_tolerances
                )
                return is_match, confidence, "traditional"

        except Exception as e:
            logger.error(f"混合颜色检测失败: {e}")
            return False, 0.0, "error"

    def _traditional_detection(self, target_color, pixels, algorithm, tolerance, hsv_tolerances):
        """传统颜色检测方法"""
        try:
            if not pixels:
                return False, 0.0

            # 计算平均颜色
            avg_color = np.mean(pixels, axis=0).astype(int)

            if algorithm == "RGB":
                # RGB欧氏距离
                distance = np.linalg.norm(np.array(target_color) - avg_color)
                is_match = distance <= tolerance
                confidence = max(0, 1 - distance / tolerance) if is_match else 0
            else:
                # HSV比较
                import colorsys

                # 转换为HSV
                target_hsv = colorsys.rgb_to_hsv(*[c/255 for c in target_color])
                current_hsv = colorsys.rgb_to_hsv(*[c/255 for c in avg_color])

                # 计算HSV差异
                h_diff = min(abs(target_hsv[0] - current_hsv[0]),
                           1 - abs(target_hsv[0] - current_hsv[0]))
                s_diff = abs(target_hsv[1] - current_hsv[1])
                v_diff = abs(target_hsv[2] - current_hsv[2])

                # 检查是否在容差范围内
                h_tol = hsv_tolerances["h"] / 360
                s_tol = hsv_tolerances["s"] / 100
                v_tol = hsv_tolerances["v"] / 100

                is_match = (h_diff <= h_tol and s_diff <= s_tol and v_diff <= v_tol)

                # 计算置信度
                if is_match:
                    h_conf = 1 - h_diff / h_tol
                    s_conf = 1 - s_diff / s_tol
                    v_conf = 1 - v_diff / v_tol
                    confidence = (h_conf + s_conf + v_conf) / 3
                else:
                    confidence = 0

            return is_match, confidence

        except Exception as e:
            logger.error(f"传统颜色检测失败: {e}")
            return False, 0.0

# 性能测试函数
def benchmark_color_detection():
    """性能基准测试"""
    print("🔬 颜色检测方法性能对比测试")
    print("=" * 50)

    # 创建测试数据
    test_sizes = [25, 100, 400, 1600]  # 不同的检测区域大小
    test_iterations = 100

    kmeans_detector = KMeansColorDetector()
    hybrid_detector = HybridColorDetector()

    results = {}

    for size in test_sizes:
        print(f"\n测试区域大小: {size} 像素")

        # 生成测试图像数据
        np.random.seed(42)
        test_pixels = np.random.randint(0, 256, (size, 3))
        target_color = [128, 128, 128]  # 灰色目标

        # 测试K-Means方法
        kmeans_times = []
        for _ in range(test_iterations):
            start_time = time.time()
            kmeans_detector.is_color_match_kmeans(target_color, test_pixels)
            kmeans_times.append(time.time() - start_time)

        # 测试传统方法
        traditional_times = []
        for _ in range(test_iterations):
            start_time = time.time()
            hybrid_detector._traditional_detection(
                target_color, test_pixels, "RGB", 30, None
            )
            traditional_times.append(time.time() - start_time)

        # 记录结果
        results[size] = {
            "kmeans_avg": np.mean(kmeans_times) * 1000,  # 转换为毫秒
            "traditional_avg": np.mean(traditional_times) * 1000,
            "speedup": np.mean(traditional_times) / np.mean(kmeans_times)
        }

        print(f"  K-Means平均耗时: {results[size]['kmeans_avg']:.2f}ms")
        print(f"  传统方法平均耗时: {results[size]['traditional_avg']:.2f}ms")
        print(f"  速度比: {results[size]['speedup']:.2f}x")

    return results

if __name__ == "__main__":
    # 运行性能测试
    benchmark_color_detection()
